import gradio as gr
import numpy as np
from PIL import Image
import pickle

# 加载模型
with open('best_knn_model.pkl', 'rb') as file:
    knn_model = pickle.load(file)

def preprocess(image):
    """将输入的 PIL 图像转换为模型需要的格式"""
    # 确保图像是灰度
    if image.mode != 'L':
        image = image.convert('L')
    # 调整图像大小到 28x28 像素
    image = image.resize((8, 8), Image.ANTIALIAS)
    # 将图像转换为 NumPy 数组并归一化
    image_array = np.array(image, dtype=np.float32) / 255.0
    # 展平数组
    image_array = image_array.flatten()
    return image_array.reshape(1, -1)

def predict(image):
    """使用 KNN 模型进行预测"""
    processed_image = preprocess(image)
    prediction = knn_model.predict(processed_image)
    return prediction[0]

# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(label="Draw your digit", type="pil"),
    outputs="label",
    title="Handwritten Digit Recognition",
    description="Draw a digit in the box and we'll tell you what it is!"
)

# 启动应用
iface.launch()